import os
import re
from multiprocessing import Process
from threading import Thread, Lock

import requests
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad

LOCK = Lock()


def decry_aes(key, iv, download_path, ts_name, data):
    if len(key):  # AES 解密
        cryptor = AES.new(key, AES.MODE_CBC, iv)
        try:
            if len(data) % 16:
                print('需要填充！！！')
                BLOCK_SIZE = 16
                with open(f'{ts_name}.txt', 'a', encoding='utf8')as f:
                    f.write(f'{download_path}-{ts_name}填充\n')
                with open(os.path.join(download_path, ts_name), 'ab') as f:
                    f.write(cryptor.decrypt(pad(data, BLOCK_SIZE)))
                print(f'{ts_name}写入成功')
            else:
                with open(os.path.join(download_path, ts_name), 'ab') as f:
                    f.write(cryptor.decrypt(data))
                print(f'{ts_name}写入成功')
        except Exception as f:
            print(f'{ts_name}写入失败', f)
            with open(f'{ts_name}.txt')as f:
                f.write(f'{download_path}-{ts_name}失败\n')


def download(url, title, chapter, course_id):
    try:
        iv = re.findall('video/(.*?)/', url)[0].encode('utf8')  # 偏移量
        print('iv:', iv)
    except:
        print("取iv fail")
        return
    LOCK.acquire()  # 加锁  保护os执行路径安全
    download_path = os.getcwd() + f"\\{course_id}_mp4"
    LOCK.release()  # 解锁
    if not os.path.exists(download_path):
        try:
            os.mkdir(download_path)
        except FileExistsError:
            pass
    # 新建章节文件夹
    download_path = os.path.join(download_path, chapter)

    if not os.path.exists(download_path):

        try:
            os.mkdir(download_path)
        except FileExistsError:
            pass

    download_path = os.path.join(download_path, title)

    if not os.path.exists(download_path):

        try:
            os.mkdir(download_path)
        except FileExistsError:
            pass
    if os.path.exists(os.path.join(download_path, f"{title}.mp4")):
        print("mp4文件已存在")
        return
    all_data = requests.get(url, timeout=9).text  # 获取M3U8文件内容
    if "#EXTM3U" not in all_data:
        raise BaseException("非M3U8的链接")

    data_line = all_data.split("\n")

    unknow = True
    key = ""  # # 秘钥
    fail_li = []  # ts_url出错列表
    for index, line in enumerate(data_line):
        if "#EXT-X-KEY" in line:  # 找解密Key

            index_l = line.find("URI")
            index_r = line.rfind('"')
            key_path = line[index_l:index_r].split('"')[1]

            key_url = url.rsplit("/", 1)[0] + "/" + key_path  # 拼出key解密密钥URL
            try:
                res = requests.get(key_url, timeout=9)  # 拿到key
                key = res.content
            except Exception as f:
                print("访问key_url fail", f)
                return
            else:
                print("key：", key)

        if "EXTINF" in line:  # 找ts地址并下载
            unknow = False
            ts_url = url.rsplit("/", 1)[0] + "/" + data_line[index + 1]  # 拼出ts片段的URL
            ts_name = data_line[index + 1].rsplit("/", 1)[-1]
            if os.path.exists(os.path.join(download_path, ts_name)):
                print('文件已存在!!')
                continue
            try:
                data = requests.get(ts_url, timeout=9).content
            except Exception as f:
                print(f'[{ts_name}] ts_url fail ', f)
                fail_li.append(ts_url)
                continue
            decry_aes(key=key, iv=iv, download_path=download_path, ts_name=ts_name, data=data)

    if fail_li:
        print('处理失败url')
        for index, ts_url in enumerate(fail_li):
            ts_name = re.findall('m3u8/(.*?).ts', ts_url)[0] + '.ts'
            msg = 1
            while msg:
                try:
                    print(f'第{msg}次尝试！！')
                    data = requests.get(ts_url, timeout=9).content
                except Exception as f:
                    print(f'[fail_{ts_name}] ts_url fail ', f)
                    msg += 1
                    continue
                else:
                    msg = 0
                    decry_aes(key=key, iv=iv, download_path=download_path, ts_name=ts_name, data=data)
    if unknow:
        raise BaseException("未找到对应的下载链接")
    else:
        print("下载完成")

    merge_file(download_path, title)


def merge_file(path, title):
    file_li = os.listdir(path)

    file_li.sort(key=lambda x: int(x[5:-3]))

    li = []
    length = len(file_li)
    if length > 400:  # 如果ts文件过多，一次合并会失败，按照400一批拆成多批合并
        li.append(file_li[:400])
        li.append(file_li[400:])
        li = [file_li[i:i+400] for i in range(0, length, 400)]
    else:
        li.append(file_li)
    for num, file_li in enumerate(li):
        file_cmd = ''
        for index, file in enumerate(file_li):

            if index != 0:
                file_cmd = file_cmd + '+' + f"{file}"
            else:
                file_cmd = f"{file}"

        cmd = f'copy /b {file_cmd} "{path}\\{title}_{num}.mp4"'  # 合并成子MP4文件语句

        LOCK.acquire()  # 加锁  保护os执行路径安全
        path1 = os.getcwd()
        os.chdir(path)
        os.system(cmd)
        cmd = f'copy /b *_?.mp4 "{path}\\{title}.mp4"'  # 将子mp4文件合并主mp4语句
        os.system(cmd)
        os.chdir(path1)
        LOCK.release()  # 解锁

    if os.path.exists(f"{path}\\{title}.mp4"):
        os.system(f'del /Q "{path}\\*.ts"')  # 删除ts文件
        os.system(f'del /Q "{path}\\*_*.mp4"')  # 删除_ 的临时MP4文件


def get_url(access_token, course_id, class_id, unit_id):
    headers = {
        'Referer': f'https://www.bjjnts.cn/study/video?course_id={course_id}&unit_id={unit_id}&class_id={class_id}',
        'Authorization': f'Bearer {access_token}',  ##登录cookie
        'Origin': 'https://www.bjjnts.cn'
    }
    res = requests.get(f'https://apif.bjjnts.cn/course-units/{unit_id}?class_id={course_id}', headers=headers,
                       timeout=9).json()
    if not res['video']:
        raise BaseException('拿m3u8 url fail')
    return res


def login():
    url = 'https://apif.bjjnts.cn/account/login'
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.45 Safari/537.36',
        'Origin': 'https://www.bjjnts.cn',

    }
    data = {"username": "15****78376", "password": "****", "type": 1}
    try:
        res = requests.post(url=url, headers=headers, data=data, timeout=9).json()

        access_token = res['access_token']
        if access_token:
            print('登录成功')
            return access_token
    except:
        raise BaseException('登录失败')


def get_course_li(course_id, class_id, access_token):
    url = f'https://apif.bjjnts.cn/courses/test-preview?course_id={course_id}&class_id={class_id}'
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.45 Safari/537.36',
        'Origin': 'https://www.bjjnts.cn',
        'Referer': f'https://www.bjjnts.cn/study?course_id={course_id}&class_id={class_id}'
    }
    try:
        res = requests.get(url=url, headers=headers, timeout=9).json()

        course_li = res['course']
    except:
        raise BaseException('取课程列表 fail')
    print(course_li)
    return course_li


def start_thread(i, access_token, course_id, class_id, chapter, index):
    unit_id = i['id']
    res = get_url(access_token, course_id=course_id, class_id=class_id, unit_id=unit_id)
    title = res['title']
    url = res['video']['urls'][0]['url']
    if url:
        print('取m3u8 url success')
    print(url)
    download(url, title, chapter + str(index), course_id)


def start_process(access_token, value, course_id, class_id, index):
    chapter = value['title']
    for i in value["units"]:
        Thread(target=start_thread, args=(i, access_token, course_id, class_id, chapter, index)).start()


def main():
    class_id = 42654  # 班级id
    course_id_li = [1946]  # 课程id列表

    access_token = login()  # 登录

    for _, course_id in enumerate(course_id_li):

        course_li = get_course_li(course_id, class_id, access_token)
        for index, value in enumerate(course_li):
            Process(target=start_process, args=(access_token, value, course_id, class_id, index)).start()


if __name__ == '__main__':
    main()
